library(dplyr)
library(tidyr)
library(ggplot2)
library(viridis)
library(ggpubr)
library(cowplot)
library(gridExtra)
library(stringr)


if (!file.exists("output")) {
  dir.create("output")
}

# load dara
# CN clusters generated in Figure_7c.R
CN_clusters <- read.csv("output/CN_clusters.csv")

cn_profile <- read.csv("ExtendedDataFigure_8b_data.csv") %>% left_join(CN_clusters, by = "donor_id")

# plot ploidy and CNV burden
for (cnv in c("ploidy", "CNV_burden")) {
  p <- ggplot(cn_profile, aes(x = HNC_clusters, y = cn_profile[, cnv], fill = HNC_clusters)) +
    geom_violin(aes(color = HNC_clusters), width = 1, alpha = .8, linewidth = .3) +
    geom_boxplot(color = "#585858", alpha = 0.15, linewidth = .3, staplewidth = .2, outlier.shape = NA) +
    geom_jitter(fill = "#585858", size = .8, stroke = 0, alpha = 0.25, width = .1, height = .1) +
    labs(title = str_replace(cnv, "_", " "), y = ifelse(cnv == "ploidy", "ploidy", "number of events"), x = "HNC clusters") +
    ylim(NA, max(cn_profile[, cnv]) * 1.1) +
    scale_fill_viridis(discrete = TRUE, end = .8) +
    scale_color_viridis(discrete = TRUE, end = .8) +
    theme_minimal() +
    theme(
      legend.position = "none",
      axis.text.x = element_text(angle = 45, hjust = 1),
      plot.title = element_text(face = "bold", size = 12, hjust = 0.5),
      plot.subtitle = element_text(size = 10),
      axis.text = element_text(size = 12),
      axis.line = element_line(size = .3,color = "#404040"),
      axis.ticks = element_line(size = .3,color = "#404040"),
      panel.grid = element_blank()
    ) +
    stat_compare_means(label.x = length(unique(cn_profile$HNC_clusters)) / 2 + 0.5, hjust = 0.5, size = 3) #+
  print(p)
  ggsave(paste0("output/Extended_Data_Figure_8b_", cnv, Sys.Date(), ".pdf"),
    device = "pdf",
    width = 3.5, height = 4, dpi = 700
  )
}

# plot CNV gains, losses and NLOH
for (cnv in c("Gain", "Loss", "NLOH")) {
  p <- ggplot(cn_profile, aes(x = HNC_clusters, y = cn_profile[, cnv], fill = HNC_clusters)) +
    geom_violin(aes(color = HNC_clusters), width = 1, alpha = .8, linewidth = .3) +
    geom_boxplot(color = "#585858", alpha = 0.15, linewidth = .3, staplewidth = .2, outlier.shape = NA) +
    geom_jitter(fill = "#585858", size = .8, stroke = 0, alpha = 0.25, width = .1, height = .1) +
    labs(title = paste("CN", cnv), y = "number of events", x = "HNC clusters") +
    ylim(NA, max(cn_profile[, cnv]) * 1.1) +
    scale_fill_viridis(discrete = TRUE, end = .8) +
    scale_color_viridis(discrete = TRUE, end = .8) +
    theme_minimal() +
    theme(
      legend.position = "none",
      axis.text.x = element_text(angle = 45, hjust = 1),
      plot.title = element_text(face = "bold", size = 12, hjust = 0.5),
      plot.subtitle = element_text(size = 10),
      axis.text = element_text(size = 12),
      axis.line = element_line(size = .3,color = "#404040"),
      axis.ticks = element_line(size = .3,color = "#404040"),
      panel.grid = element_blank()
    ) +
    stat_compare_means(label.x = length(unique(cn_profile$HNC_clusters)) / 2 + 0.5, hjust = 0.5, size = 3) #+
  print(p)
  ggsave(paste0("output/Extended_Data_Figure_8b_", cnv, "_HNCclusters_", Sys.Date(), ".pdf"),
    device = "pdf",
    width = 3.5, height = 4, dpi = 700
  )
}
